#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：gas_pipeline_qc_partition4.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 16:25 
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
QC scheduling with device-selection encoded in GA
-------------------------------------------------
- 不考虑充气细节：同配方视为一次性充好，所有瓶同时进入QC（mix_end=0）。
- 每个组分有两台设备（只测该组分）。
- 每瓶的分组策略 + 每组分的设备选择 都在染色体里编码。
- 组大小=1 → 串行；组大小>=2 → 1-pass（同起同止，段长=max(单机时长)），中途不允许拆换。
- GA 最小化 Cmax，输出收敛曲线与甘特图（全黑字）。

Author: you
"""
import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ============================ CONFIG（可自定义） ============================ #
CONFIG = {
    # 规模
    "NUM_BOTTLES": 18,                 # 气瓶数量（工件数）
    "MAX_COMPONENTS_PER_BOTTLE": 4,    # 每瓶最大组分数（1..K）

    # 组分全集（自动为每个组分生成两台设备）
    "COMPONENT_TYPES": ["A","B","C","D","O2","CO2","CH4","H2","CO","N2"],

    # QC 单机检测时间（小时）
    "COMPONENT_TEST_TIME": {
        "A":0.22, "B":0.20, "C":0.25, "D":0.18,
        "O2":0.18, "CO2":0.22, "CH4":0.25, "H2":0.15, "CO":0.20, "N2":0.12
    },

    # QC 换型时间（同一设备上若上次签名≠这次的组件名，需要 setup）
    "QC_SETUP_TIME": 0.12,

    # GA 参数（优化：分组策略 + 设备选择）
    "GA_POP": 80,
    "GA_GENS": 180,
    "GA_PC": 0.9,
    "GA_PM_POL": 0.20,    # policy 变异率（每个体）
    "GA_PM_DEV": 0.10,    # 设备选择变异率（每瓶-每组分翻转的概率）
    "GA_ELITE": 10,

    # 绘图
    "LABEL_WRAP_CHARS": 22,
    "RANDOM_SEED": 7
}

# ============================ 数据结构 ============================ #
@dataclass(frozen=True)
class Recipe:
    components: Tuple[str, ...]   # 不考虑浓度，只要组分相同视作同配方
    def key(self): return tuple(sorted(self.components))

@dataclass
class Bottle:
    id: int
    recipe: Recipe

@dataclass
class QCTask:
    start: float
    end: float
    device: int           # 设备全局索引
    label: str
    group_key: Optional[Tuple[int,float,float,int]] = None  # (bid, st, en, group_idx) for 1-pass coloring

# ============================ 随机实例生成 ============================ #
class RandomData:
    """
    - 每个组分生成两台设备（comp -> [dev_id0, dev_id1]）
    - 每瓶从组件全集中抽 1..K 个组分
    """
    def __init__(self, cfg):
        self.cfg = cfg
        random.seed(cfg["RANDOM_SEED"]); np.random.seed(cfg["RANDOM_SEED"])

    def gen(self):
        comps_all = self.cfg["COMPONENT_TYPES"]
        # 创建设备：每个组分两台
        comp_to_devices: Dict[str, List[int]] = {}
        device_labels: List[str] = []
        device_index = 0
        for c in comps_all:
            comp_to_devices[c] = [device_index, device_index+1]
            device_labels.append(f"{c}-1")
            device_labels.append(f"{c}-2")
            device_index += 2
        D = device_index  # 总设备数 = 2 * |components|

        # 生成气瓶
        N = self.cfg["NUM_BOTTLES"]
        Kmax = self.cfg["MAX_COMPONENTS_PER_BOTTLE"]
        bottles: List[Bottle] = []
        for i in range(N):
            k = np.random.randint(1, Kmax+1)
            k = min(k, len(comps_all))
            comps = random.sample(comps_all, k)
            bottles.append(Bottle(i, Recipe(tuple(comps))))
        return bottles, comp_to_devices, device_labels

# ============================ 调度器（仅 QC） ============================ #
class Scheduler:
    """
    仅分析阶段调度：
    - 染色体包含：每瓶分组策略（policy） + 每组分设备选择（二选一）。
    - 组大小=1（串行）：在染色体指定的那台设备上排产；
    - 组大小>=2（1-pass）：每个组分用染色体指定的设备；所有设备同起同止，段长=max(该组单机时长)，中途不允许拆换。
    """
    def __init__(self, cfg, bottles: List[Bottle], comp_to_devices: Dict[str, List[int]], device_labels: List[str]):
        self.cfg = cfg
        self.bottles = bottles
        self.N = len(bottles)
        self.comp_to_devices = comp_to_devices
        self.device_labels = device_labels
        self.D = len(device_labels)

    def single_test_time(self, comp: str) -> float:
        return float(self.cfg["COMPONENT_TEST_TIME"].get(comp, 0.2))

    # policy codes: 0=ALL_TOGETHER, 1=SPLIT_LONGEST, 2=EACH_SEPARATE
    def make_groups(self, comps: List[str], policy: int) -> List[List[str]]:
        if len(comps) <= 1:
            return [comps[:]]
        t = {c: self.single_test_time(c) for c in comps}
        if policy == 0:   # ALL together
            return [comps[:]]
        if policy == 1:   # longest alone + rest together
            longest = max(comps, key=lambda x: t[x])
            rest = [c for c in comps if c != longest]
            return [[longest], rest] if rest else [[longest]]
        # EACH separate
        return [[c] for c in sorted(comps)]

    def schedule_qc_with_chrom(self,
                               policies: np.ndarray,
                               dev_bits: List[Dict[str,int]]
                               ) -> Tuple[float, List[QCTask]]:
        """
        输入：
          - policies[i] ∈ {0,1,2} 表示第 i 瓶的分组策略；
          - dev_bits[i][comp] ∈ {0,1} 表示第 i 瓶该组分选两台设备的哪一台；
        调度：
          - 串行：用指定设备；1-pass：每个组分用其指定设备；同起同止；setup 依据设备上次签名（上次组分名）。
        """
        D = self.D
        ready = [0.0]*D
        last_sig: List[Optional[str]] = [None]*D
        setup_t = self.cfg["QC_SETUP_TIME"]

        tasks: List[QCTask] = []
        # 所有瓶 mix_end=0，按 id 顺序进入 QC
        order_seq = sorted(self.bottles, key=lambda b: b.id)

        for b in order_seq:
            bid = b.id
            t_cur = 0.0
            comps = list(b.recipe.components)
            groups = self.make_groups(comps, int(policies[bid]))
            choice = dev_bits[bid]  # dict: comp -> 0/1

            for gi, group in enumerate(groups):
                if len(group) == 1:
                    comp = group[0]
                    dev_ids = self.comp_to_devices[comp]
                    pick = int(choice.get(comp, 0)) % 2
                    d = dev_ids[pick]
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                    st = max(t_cur, ready[d] + stp)
                    pt = self.single_test_time(comp)
                    en = st + pt
                    label = f"S{bid} serial {self.device_labels[d]}({comp})"
                    tasks.append(QCTask(st, en, d, label, None))
                    ready[d] = en
                    last_sig[d] = comp
                    t_cur = en
                    continue

                # 1-pass：每个组分在染色体选择的设备上同时开始/结束
                chosen: List[Tuple[int, str, float]] = []  # (device, comp, rdy_with_setup)
                earliest = t_cur
                for comp in group:
                    dev_ids = self.comp_to_devices[comp]
                    pick = int(choice.get(comp, 0)) % 2
                    d = dev_ids[pick]
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==comp) else setup_t
                    rdy = ready[d] + stp
                    chosen.append((d, comp, rdy))
                    earliest = max(earliest, rdy)

                seg_len = max(self.single_test_time(c) for (_, c, _) in chosen)
                st = earliest; en = st + seg_len
                used = "+".join(self.device_labels[d] for (d,_,_) in chosen)
                comps_str = ",".join(sorted(group))
                label = f"S{bid} 1-pass[{used}] comps={comps_str} max={seg_len:.2f}h"
                group_key = (bid, round(st,6), round(en,6), gi)
                for (d, comp, _) in chosen:
                    tasks.append(QCTask(st, en, d, label, group_key))
                    ready[d] = en
                    last_sig[d] = comp
                t_cur = en

        Cmax = max((t.end for t in tasks), default=0.0)
        return Cmax, tasks


# ============================ GA（分组策略 + 设备选择） ============================ #
class Chromosome:
    """
    policies: np.ndarray shape (N,), values in {0,1,2}
    dev_bits: List[Dict[str,int]], len=N, each maps comp-> {0,1}
    """
    def __init__(self, policies: np.ndarray, dev_bits: List[Dict[str,int]]):
        self.policies = policies
        self.dev_bits = dev_bits

    def copy(self):
        return Chromosome(self.policies.copy(), [d.copy() for d in self.dev_bits])

class GA:
    """
    染色体包含：
      - 每瓶分组策略（0/1/2）
      - 每瓶每组分的设备选择（两台中选一台，用 bit 表示）
    交叉：
      - policy：单点交叉
      - dev_bits：对两亲本公共的 (i, comp) 做随机继承；其余沿用各自
    变异：
      - policy：以 PM_POL 概率随机重置一位
      - dev_bits：以 PM_DEV 概率对每瓶每组分翻转 bit
    """
    def __init__(self, cfg, scheduler: Scheduler):
        self.cfg = cfg
        self.S  = scheduler
        self.N  = self.S.N
        random.seed(cfg["RANDOM_SEED"]); np.random.seed(cfg["RANDOM_SEED"])

    def rand_ind(self) -> Chromosome:
        policies = np.random.randint(0,3,size=self.N,dtype=int)
        dev_bits: List[Dict[str,int]] = []
        for i,b in enumerate(self.S.bottles):
            mp={}
            for c in b.recipe.components:
                mp[c] = np.random.randint(0,2)
            dev_bits.append(mp)
        return Chromosome(policies, dev_bits)

    def crossover(self, A: Chromosome, B: Chromosome):
        # policy 单点
        n = self.N
        cut = np.random.randint(1, n)
        p1 = np.concatenate([A.policies[:cut], B.policies[cut:]])
        p2 = np.concatenate([B.policies[:cut], A.policies[cut:]])

        # dev_bits：对每瓶，对“两个亲本都包含的组分键”随机继承；只在交集上重组，避免丢键
        d1=[]; d2=[]
        for i in range(n):
            mA=A.dev_bits[i]; mB=B.dev_bits[i]
            keys = set(mA.keys()) | set(mB.keys())
            child1={}; child2={}
            for k in keys:
                if (k in mA) and (k in mB):
                    if np.random.rand()<0.5:
                        child1[k]=mA[k]; child2[k]=mB[k]
                    else:
                        child1[k]=mB[k]; child2[k]=mA[k]
                elif k in mA:
                    child1[k]=mA[k]; child2[k]=mA[k]
                else:
                    child1[k]=mB[k]; child2[k]=mB[k]
            d1.append(child1); d2.append(child2)

        return Chromosome(p1,d1), Chromosome(p2,d2)

    def mutate(self, X: Chromosome):
        # policy 变异
        if np.random.rand() < self.cfg["GA_PM_POL"]:
            i=np.random.randint(self.N)
            X.policies[i]=np.random.randint(0,3)
        # dev_bits 变异（逐瓶逐组分翻转）
        pm = self.cfg["GA_PM_DEV"]
        for i in range(self.N):
            for k in list(X.dev_bits[i].keys()):
                if np.random.rand()<pm:
                    X.dev_bits[i][k] = 1 - int(X.dev_bits[i][k])
        return X

    def fitness(self, X: Chromosome):
        Cmax, _ = self.S.schedule_qc_with_chrom(X.policies, X.dev_bits)
        return Cmax

    def run(self):
        pop=self.cfg["GA_POP"]; gens=self.cfg["GA_GENS"]
        pc=self.cfg["GA_PC"]; elite=self.cfg["GA_ELITE"]

        pool=[]
        for _ in range(pop):
            ind=self.rand_ind()
            f=self.fitness(ind)
            pool.append((ind,f))
        best=min(pool, key=lambda x:x[1]); hist=[best[1]]

        for g in range(1, gens+1):
            pool.sort(key=lambda x:x[1])
            elites=pool[:elite]
            new=[(e[0].copy(), e[1]) for e in elites]
            while len(new)<pop:
                A=random.choice(pool)[0]
                B=random.choice(pool)[0]
                C1, C2 = A.copy(), B.copy()
                if np.random.rand()<pc:
                    C1, C2 = self.crossover(A, B)
                C1 = self.mutate(C1); f1=self.fitness(C1)
                new.append((C1, f1))
                if len(new)<pop:
                    C2 = self.mutate(C2); f2=self.fitness(C2)
                    new.append((C2, f2))
            pool=new
            cur=min(pool, key=lambda x:x[1])
            if cur[1] < best[1]:
                best=cur
            hist.append(best[1])
        return best[0], best[1], hist

# ============================ 绘图 ============================ #
def wrap_text(s: str, width: int) -> str:
    out=[]; line=""
    for ch in s:
        line += ch
        if len(line) >= width and ch == ' ':
            out.append(line.rstrip()); line=""
    if line: out.append(line)
    return "\n".join(out)

def plot_convergence(hist, title="GA Convergence (minimize Cmax)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return
    plt.figure(figsize=(6,4))
    plt.plot(hist, linewidth=2, color='black')
    plt.xlabel("Generation", color='black'); plt.ylabel("Best Cmax", color='black')
    plt.title(title, color='black'); plt.grid(True, alpha=0.3, color='black')
    ax=plt.gca()
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    for spine in ax.spines.values(): spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()

def plot_gantt(qc_tasks: List[QCTask], device_labels: List[str], cfg,
               title="QC Gantt (device choices encoded; 1-pass same [st,en])"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return
    D = len(device_labels)
    row_idx = {d: d for d in range(D)}

    # 为1-pass组统一颜色
    from itertools import cycle
    palette = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
    grp_color: Dict[Tuple[int,float,float,int], str] = {}
    for t in qc_tasks:
        if t.group_key and t.group_key not in grp_color:
            grp_color[t.group_key] = next(palette)

    fig_h = max(4.5, 0.6*D + 2)
    fig, ax = plt.subplots(figsize=(12, fig_h))

    for t in qc_tasks:
        y=row_idx[t.device]
        color = grp_color.get(t.group_key, None)
        ax.barh(y, t.end-t.start, left=t.start, height=0.8,
                color=color, edgecolor='black', linewidth=0.6)
        ax.text((t.start+t.end)/2, y, wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    ax.set_yticks(list(range(D)))
    ax.set_yticklabels([f"E{d+1} ({device_labels[d]})" for d in range(D)], color='black')
    ax.set_xlabel("Time (h)", color='black'); ax.set_title(title, color='black')
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    ax.grid(axis='x', alpha=0.3, color='black')
    for spine in ax.spines.values(): spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()

# ============================ 主程序 ============================ #
def main():
    random.seed(CONFIG["RANDOM_SEED"]); np.random.seed(CONFIG["RANDOM_SEED"])

    # 1) 实例（设备两台/组分；瓶的组分随机）
    bottles, comp_to_devices, device_labels = RandomData(CONFIG).gen()

    # 2) 调度器
    sch = Scheduler(CONFIG, bottles, comp_to_devices, device_labels)

    # 3) GA：优化“分组策略 + 设备选择”
    ga = GA(CONFIG, sch)
    best_ind, best_cmax, hist = ga.run()

    # 4) 用最优个体排程，输出结果并绘图
    Cmax, qc_tasks = sch.schedule_qc_with_chrom(best_ind.policies, best_ind.dev_bits)

    print(f"Bottles={CONFIG['NUM_BOTTLES']}, Components={len(CONFIG['COMPONENT_TYPES'])}, Devices={len(device_labels)} (2 per component)")
    print("Best policy per bottle (0=ALL,1=SPLIT_LONGEST,2=EACH):")
    print(best_ind.policies.tolist())
    print("Best device bits per bottle (component -> 0/1):")
    for i,b in enumerate(bottles):
        print(f"S{i}: {best_ind.dev_bits[i]}")
    print(f"Best Cmax: {Cmax:.3f} h")

    if HAS_MPL:
        plot_convergence(hist)
        plot_gantt(qc_tasks, device_labels, CONFIG)
    else:
        print("(Install matplotlib to view plots.)")

if __name__ == "__main__":
    main()
